package org.tigr.htc.server.sge;

import java.io.File;
import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.util.Arrays;
import java.util.Calendar;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.apache.log4j.Logger;
import org.tigr.antware.shared.command.events.CommandInterrupt;
import org.tigr.antware.shared.util.StringUtils;
import org.tigr.htc.cmd.Command;
import org.tigr.htc.common.HTCConfig;
import org.tigr.antware.shared.util.RunSystemCommand;
import org.tigr.htc.server.IJobMonitor;

/**
 * The <b>SgeJobMonitor</b> class - is a utility class that is used to monitor one or
 * more of jobs submitted to Sun Grid Engine (SGE). This class monitors all the jobs that have
 * been registered with this class for monitoring. When the runner submits a job it 
 * registers the job id and when the job is done the job id is removed from the monitored list.
 * This class periodically monitors the status of the jobs and caches the status
 * at the specified frequency.
 *
 */
public class SgeJobMonitor implements IJobMonitor {
    static Logger log = Logger.getLogger(SgeJobMonitor.class);

    /**
     * The variable <code>frequency</code> holds the frequency in milliseconds at which 
     * the job status are monitored.
     */
    private static int frequency = 5000;
    
    /**
     * The variable <code>jobs</code> holds the map of registered jobs to monitor and the
     * underlying command objects. This map is organized by the Grid jobID
     */
    private Map jobsByGridID = Collections.synchronizedMap(new HashMap());
    
    /**
     * The variable <code>statuses</code> holds the map of registered jobs and their 
     * status. This map is organized by the Grid jobID
     */
    private Map statuses = Collections.synchronizedMap(new HashMap());
    
    /**
     * The variable <code>active</code> us used to decide if the monitor thread should die.
     * If the active flag is true the thread keeps monitoring for the change in job status
     */
    private boolean active = true;

    // Command name to check the completed status 
    final static String htc_complete = HTCConfig.getProperty("sge_complete");

    // Command name to check the running status
    final static String htc_running = HTCConfig.getProperty("sge_running");
    
    final static String jobIDPattern = HTCConfig.getProperty("sge_job_id");

    /** 
     * The method <code>setActive</code> sets the value for active. If the flag is set to
     * false, the monitor stops monitoring the status.
     * 
     * @param active A <code>boolean</code> value for the field active.
     */
    public void setActive(boolean active) {
        this.active = active;
    }
    
    /**
     * The method <code>register</code> registers the job with the specified GRID ID to list
     * of jobs being monitored by this job monitor.
     * 
     * @param pGridJobID a <code>int</code> representing the underlying GRID jobID
     * @param pCmd a <code>Command</code>
     */
    public void register(String pGridJobID, Command pCmd) {
    	// As the map may be in the process of being used to get the status
        // synchronize on the vector
        synchronized (jobsByGridID) {
            log.debug("Registering grid job ID: " + pGridJobID + " HTC ID: " 
                    + pCmd.getID());
            statuses.put(pGridJobID, new SgeJobStatus());
            jobsByGridID.put(pGridJobID, pCmd);
            log.debug("Registered job ID: " + pGridJobID);
        }
    }

    /**
     * The method <code>run</code> will launch the thread that monitors the jobs at the 
     * specified frequency
     *
     */
    public void run() {
    	// Loop till the monitor is active
        while (active) {
            try {
                synchronized (jobsByGridID) {
                    if (jobsByGridID.size() > 0) {
                        log.debug("Attempting to read the status of jobs.");
                        readStatus();
                        log.debug("Finished reading status.");
                    }
                }
            } catch (Throwable e) {
				log.error("Error monitoring jobs", e);
            }
            
            try {
				Thread.sleep(frequency);
			} catch (InterruptedException e) {
				// ignored
			}
            log.debug("Finished sleeping. Will attempt to read status.");
		}
        log.info("Monitoring complete!");
    }

    /**
     * The method <code>readStatus</code> reads and sets the status for the all the jobs
     * in the hash map.
     * The method uses <code>qstat</code> to see if the job is idle or running, and
     * the <code>qacct</code> command to get the successful state and return value.
     * If the status of any of the jobs has changed since the last try, the method fires
     * of the appropriate events.
     * 
     * @throws Exception
     */
    public void readStatus() throws Throwable {
        log.debug("In readStatus.");
        SgeJobStatus status = null;
        String strStartTime = null;
        String strEndTime = null;
        SimpleDateFormat qacct_df = new SimpleDateFormat("EEE MMM dd HH:mm:ss yyyy");
        SimpleDateFormat qstat_df = new SimpleDateFormat("MM/dd/yyyy HH:mm:ss");
        
        // Holds the list of jobIDs whose status has been checked already
        HashSet statusChecked = new HashSet();
        
        // Execute the 'qstat' command to get the status of active jobs
        String sgeCmd = htc_running;
        RunSystemCommand run = new RunSystemCommand();
        run.runCommand(sgeCmd);
        
        // If the command completes without any errors, then grab the output
        // and process all the jobs of interest. This will give the status of all the idling, held and
        // running jobs. The next step checks for all the completed jobs
        if (run.getExitValue() == 0) {
            String output = null;
            String host = null;
            String [] lines = run.getOutString().split("\n");
            for (int i = 0; i < lines.length; i++) {
                output = lines[i];
                log.debug("Processing status line: '" + output + "'");
                if (output.length() == 0) {
                	continue; // If blank line go to next line
                }
                
                // If this line begins with jobID then this is the header line so ignore
                if (output.startsWith("job")) {
                	continue;
                }
                
                // If this line begins with '-----' then this is the separator line
                // so ignore
                if (output.startsWith("--")) {
                	continue;
                }
                
                // Else strip the leading white space and extract the job ID. If this 
                // job ID is of interest then process, otherwise skip line
                output = output.trim();
                String curGridJobID = output.substring(0, output.indexOf(" "));
                if (!jobsByGridID.containsKey(curGridJobID)) {
                	continue;
                }
                
                // Else this line corresponds to on of the jobs of interest so check to see if the 
                // job status has changed. It it has not changed go to the next line
                status = (SgeJobStatus) statuses.get(curGridJobID);
                if (! status.isStatusLineChanged(output)) {
					log.debug("Status line has not changed, so ignoring line.");
                    
                    // Add this job ID to the hash set that keeps track of checked jobs
                    // so that it is not checked again through qacct.
                    statusChecked.add(curGridJobID);
                    continue;
				}
				
                // Add this job ID to the hash set that keeps track of checked jobs
                // so that it is not checked again through qacct.
                statusChecked.add(curGridJobID);
                
                // As the status has changed reset the old status and modify the values based
                // on the new status read for this job ID
                status.reset();
                String tokens [] = output.split("\\s+");
                String state = tokens[4];
                if (state.equals("r")) {
                    status.setRunning(true);
                }
                else if (state.startsWith("h")) {
                    status.setHeld(true);
                }
                else if (state.equals("t") && (tokens.length > 7)) {
                	// If the number of tokens is greater than 7 then a host name is
                	// known so use that
                    status.setRunning(true);
                    host = tokens[7];
                    status.setHost(host);
                    
                } else {
					// If none of the above conditions is met then the job is
					// idling, so the state accordingly.
                    status.setIdle(true);
                }
                
                // If the job status is now running then send the task started event
                if (status.isRunning()) {                   
                    Command cmd = (Command) jobsByGridID.get(curGridJobID);
                    
                    // Get the actual job start time
                    strStartTime = tokens[5] + " " + tokens[6];
                    Date eventDate = qstat_df.parse(strStartTime);
					cmd.taskStarted("Command with ID: " + cmd.getID() + " started", 
                            eventDate, host, 0, 0);
				}
                
                // If an entry was found then return with the status
                log.debug("Status of job with ID: " + curGridJobID + " is " + status);
            }
        } else {
            // If there was an error executing the qstat command then throw an exception
            throw new Exception("Could not check the status for jobs.\n" 
                    + run.getErrString());
        }
        
        // After all the active jobs have been completed, get the status of jobs that
        // might be complete by using the 'qacct' command.
        // First obtain a lock on the status objects as items will be removed if 
        // jobs finish
        synchronized (statuses) {
            for (Iterator iter = statuses.keySet().iterator(); iter.hasNext();) {
                String curGridJobID = (String) iter.next();
                // Check to see if this grid job status was already checked in the
                // previous step in which case ignore it
                if (statusChecked.contains(curGridJobID)) {
                    continue;
                }
                status = (SgeJobStatus) statuses.get(curGridJobID);

                // See if the job status is already complete in which case go to next 
                // job
                if (status.isComplete()) {
                    continue;
                }

                // Execute the command to get the status of the current job
                sgeCmd = htc_complete + " -j " + curGridJobID;
                run = new RunSystemCommand();
                run.runCommand(sgeCmd);

                // If the return value is zero that means that this job is complete and
                // therefore check the completed status and return value.
                String output = null;
                
                if (run.getExitValue() == 0) {
                    String host = null;
                    boolean failed = true;
                    int returnValue = 0;

                    String[] lines = run.getOutString().split("\n");
                    for (int i = 0; i < lines.length; i++) {
                        output = lines[i];
                        log.debug("Status line: " + output);

                        // See if the line starts with 'failed' string to get the failed 
                        // status
                        if (output.startsWith("failed")) {
                            String tokens[] = output.split("\\s+");
                            String state = tokens[1];
                            if (tokens[1].equals("1")) {
                                failed = true;
                            } else {
                                failed = false;
                            }
                        }

                        // See if the line starts with 'exit_status' string to get the exit 
                        // value
                        if (output.startsWith("exit_status")) {
                            String tokens[] = output.split("\\s+");
                            String exit = tokens[1];
                            returnValue = Integer.parseInt(exit);
                        }

                        // See if the line starts with 'hostname' string if so get the 
                        // hostname
                        if (output.startsWith("hostname")) {
                            String tokens[] = output.split("\\s+");
                            host = tokens[1];
                        }

                        // Get the start time in case it was not read before
                        if (output.startsWith("start_time")) {
                            strStartTime = output.substring(output.indexOf(" ") + 1).trim();
                        }

                        // Get the end time
                        if (output.startsWith("end_time")) {
                            strEndTime = output.substring(output.indexOf(" ") + 1).trim();
                        }
                    }

                    // Check the status of the job to see if it is idle. This can happen
                    // with short lived jobs where before the monitor can register them
                    // as running they will have finished. So in such situations fire a
                    // start event before firing the finish event
                    if (status.isIdle()) {
                        log.debug("As job started was not fired, it is being fired"
                                + " before finish.");
                        Command cmd = (Command) jobsByGridID.get(curGridJobID);
                        Date startDate = Calendar.getInstance().getTime();
                        try {
                            startDate = qacct_df.parse(strStartTime);
                        } catch (ParseException e) {
                            log.warn("Error parsing start time from string '" + strStartTime + "'", e);
                        }
                        cmd.taskStarted(
                            "Command with ID: " + cmd.getID() + " started", startDate, host, 0, 0);
                    }

                    // Set the status of the job based on the failed flag and return value
                    status.setReturnValue(returnValue);
                    status.setSuccessful(failed);

                    log.debug(
                        "Job with ID: "
                            + curGridJobID
                            + " is complete. Return value: "
                            + returnValue
                            + " failed: "
                            + failed
                            + ". Will fire finish event and remove job from monitor list");

                    // At this point the job status is complete, so send the task finished 
                    // event
                    Date endDate = Calendar.getInstance().getTime();
                    try {
                        endDate = qacct_df.parse(strEndTime);
                    } catch (ParseException e) {
                        log.warn("Error parsing end time from string '" + strEndTime + "'", e);
                    }
                    Command cmd = (Command) jobsByGridID.get(curGridJobID);
                    cmd.taskFinished(
                        "Command with ID: " + cmd.getID() + " finished",
                        endDate, status.getReturnValue(), 0);

                    // If the user has sought a notification, then send the notification
                    if (cmd.getNotifyScript() != null) {
                        notify(cmd, "Command with ID: " + cmd.getID() + " finished",
                            status.getReturnValue());
                    }

                    // If the events have been fired then remove the job from monitoring
                    log.debug(
                        "Removing job with grid ID: " + curGridJobID + " from monitoring hash.");
                    iter.remove();
                    jobsByGridID.remove(curGridJobID);
                } else {
                    String error = run.getErrString();
                    log.debug("Status line: " + error);

                    // If the output says that there is nothing known about this then this
                    // is not an error so do nothing
                    if (error.indexOf("Error: job ID " + curGridJobID + " not found") >= 0) {
                        log.debug("The job with ID: " + curGridJobID + " is not complete.");
                    } else {
                        log.warn("Could not check the status for job with ID: "
                                + curGridJobID + "\n" + error);
                    }
                } // End of failure of executing qacct command
            }
        }
		
        log.debug("Leaving readStatus.");
    }
    
    /**
     * The method <code>notify</code> is used to send a notification using the notify
     * script.
     * 
     * @param cmd a <code>Command</code> object which caused this notification
     * @param message a <code>String</code> representing the message to be sent
     * @param returnValue a <code>int</code> representing the return value
     */
    public void notify(Command cmd, String message, int returnValue) {
        log.debug("In notify.");

        if (cmd.getNotifyScript() == null) {
            log.debug("No notification script set.");
            return;
        }

        String email = cmd.getEmail();
        if (email == null) {
            email = cmd.getUserName();
        }

        File scriptFile = new File(cmd.getNotifyScript());
        String cid = "--requestID=" + cmd.getID();
        String returnVal = "--returnValue=" + returnValue;
        String lmessage = "--message=" + message;
        String emailArg = "--email=" + email;
        
        if (scriptFile.exists()) {
            log.debug("Notify script file exists");
            try {
                String tmp = scriptFile.getAbsolutePath();
                
                if (log.isDebugEnabled()) {
                    log.debug("Notification command is " + tmp);
                    List args = Arrays.asList(new String[] { cid, returnVal, lmessage, emailArg, " -- " });
                    log.debug("Arguments: " + StringUtils.joinList(args, " "));
            	}
                RunSystemCommand run = new RunSystemCommand();
                String[] cmdLine = new String[] {
                        scriptFile.getAbsolutePath(),
                        cid,
                        returnVal,
                        lmessage,
                        emailArg };
                int exitVal = run.runCommand(cmdLine);
                log.debug("Notification exit value: " + exitVal);
            } catch (Exception e) {
                log.debug("Exception running notify script: " + e.getMessage(), e);
            } catch (Throwable e) {
            	log.debug("Exception running notify script: " + e.getMessage(), e);
			}
        } else {
            log.warn("Notification command was set but didn't exist: " + scriptFile);
        }
    }

    /* (non-Javadoc)
     * @see org.tigr.htc.server.IJobMonitor#stop(org.tigr.htc.cmd.Command)
     */
    public boolean stop(Command cmd) throws Exception {
        log.debug("Stopping Command with ID: " + cmd.getID() + " and grid ID: " +
                cmd.getGridJobID());
        boolean result = false;
        if ((result = SgeUtils.remove(cmd.getGridJobID()))) {
            synchronized (jobsByGridID) {
                // Remove this Command from jobs and status
                jobsByGridID.remove(cmd.getGridJobID());
                statuses.remove(cmd.getGridJobID());
            } // Fire task stopped
            cmd.taskInterrupted("Task interrupted", CommandInterrupt.USER_KILL, 0);
        }
        log.debug("Stopped Command: " + result);
        return result;
    }

    /* (non-Javadoc)
     * @see org.tigr.htc.server.IJobMonitor#hold(org.tigr.htc.cmd.Command)
     */
    public boolean hold(Command cmd) throws Exception {
        log.debug("Pausing Command with ID: " + cmd.getID() + " grid JobID: " +
                cmd.getGridJobID());
        return SgeUtils.hold(cmd.getGridJobID());
    }

    /* (non-Javadoc)
     * @see org.tigr.htc.server.IJobMonitor#resume(org.tigr.htc.cmd.Command)
     */
    public boolean resume(Command cmd) throws Exception {
        log.debug("Resuming Command with ID: " + cmd.getID() + " and grid ID: " +
                cmd.getGridJobID());
        return SgeUtils.release(cmd.getGridJobID());
    }
}
